import base64
import hashlib
import random
import string
import struct
import socket
import time
import urllib.parse
from Crypto.Cipher import AES
import xml.etree.ElementTree as ET

class WXBizMsgCrypt:
    def __init__(self, token, encoding_aes_key, corp_id):
        self.token = token
        # 注意：企业微信要求AESKey = Base64_Decode(EncodingAESKey + "=")
        self.key = base64.b64decode(encoding_aes_key + '=')
        self.corp_id = corp_id

    def get_signature(self, timestamp, nonce, msg_encrypt):
        """生成签名 - 按照企业微信官方文档"""
        # 按照字典序排序：token、timestamp、nonce、msg_encrypt
        sortlist = [self.token, timestamp, nonce, msg_encrypt]
        sortlist.sort()
        # 拼接字符串
        sort_str = ''.join(sortlist)
        # SHA1计算
        sha = hashlib.sha1()
        sha.update(sort_str.encode('utf-8'))
        return sha.hexdigest()

    def pkcs7_pad(self, text):
        """PKCS#7补位"""
        block_size = 32
        pad_amount = block_size - (len(text) % block_size)
        pad = chr(pad_amount)
        return text + pad * pad_amount

    def pkcs7_unpad(self, text):
        """去除PKCS#7补位"""
        pad = ord(text[-1])
        return text[:-pad]

    def encrypt(self, text, nonce=None, timestamp=None):
        """加密消息"""
        if not nonce:
            nonce = ''.join(random.sample(string.ascii_letters + string.digits, 16))
        if not timestamp:
            timestamp = str(int(time.time()))
        
        # 生成16字节随机字符串
        random_str = ''.join(random.sample(string.ascii_letters + string.digits, 16))
        
        # 消息长度（4字节，使用大端序）
        msg_len = struct.pack('>I', len(text))
        
        # 拼接：random(16字节) + msg_len(4字节) + msg + corp_id
        content = random_str.encode('utf-8') + msg_len + text.encode('utf-8') + self.corp_id.encode('utf-8')
        
        # PKCS#7补位
        content = self.pkcs7_pad(content.decode('latin1'))
        
        # AES加密
        iv = self.key[:16]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        encrypted = cipher.encrypt(content.encode('utf-8'))
        
        return base64.b64encode(encrypted).decode()

    def decrypt(self, encrypted):
        """解密消息 - 按照企业微信官方文档"""
        try:
            # Base64解码
            aes_msg = base64.b64decode(encrypted)
            
            # AES解密
            cipher = AES.new(self.key, AES.MODE_CBC, self.key[:16])
            rand_msg = cipher.decrypt(aes_msg)
            
            # 去除PKCS#7补位
            pad = rand_msg[-1]
            rand_msg = rand_msg[:-pad]
            
            # 去掉前16个随机字节
            content = rand_msg[16:]
            
            # 取出4字节的msg_len（注意字节序）
            msg_len_bytes = content[:4]
            msg_len = struct.unpack('>I', msg_len_bytes)[0]  # 使用大端序
            
            # 截取msg_len长度的msg
            msg = content[4:4+msg_len]
            
            # 剩余字节为receiveid
            receiveid = content[4+msg_len:]
            
            # 验证receiveid是否为corpid
            if receiveid.decode('utf-8') != self.corp_id:
                raise ValueError(f"receiveid不匹配: {receiveid.decode('utf-8')} != {self.corp_id}")
            
            return msg.decode('utf-8')
            
        except Exception as e:
            print(f"解密失败: {str(e)}")
            return None

    def verify_url(self, msg_signature, timestamp, nonce, echostr):
        """验证URL有效性"""
        try:
            # 验证签名
            signature = self.get_signature(timestamp, nonce, echostr)
            if signature != msg_signature:
                print(f"签名验证失败: {signature} != {msg_signature}")
                return None
                
            # 解密echostr
            decrypted = self.decrypt(echostr)
            return decrypted
        except Exception as e:
            print(f"URL验证失败: {str(e)}")
            return None

    def decrypt_message(self, msg_signature, timestamp, nonce, encrypted_msg):
        """解密企业微信消息"""
        try:
            # 验证签名
            signature = self.get_signature(timestamp, nonce, encrypted_msg)
            if signature != msg_signature:
                print(f"消息签名验证失败: {signature} != {msg_signature}")
                return None
                
            # 解密消息
            decrypted_xml = self.decrypt(encrypted_msg)
            return decrypted_xml
        except Exception as e:
            print(f"消息解密失败: {str(e)}")
            return None

    def parse_xml_message(self, xml_content):
        """解析XML格式的企业微信消息"""
        try:
            root = ET.fromstring(xml_content)
            
            # 提取基本信息
            msg_type = root.find('MsgType').text
            from_user = root.find('FromUserName').text
            to_user = root.find('ToUserName').text
            create_time = root.find('CreateTime').text
            msg_id = root.find('MsgId').text
            
            # 根据消息类型提取内容
            content = ""
            if msg_type == 'text':
                content = root.find('Content').text
            elif msg_type == 'image':
                pic_url = root.find('PicUrl').text
                media_id = root.find('MediaId').text
                content = f"[图片消息] {pic_url}"
            elif msg_type == 'voice':
                media_id = root.find('MediaId').text
                format = root.find('Format').text
                content = f"[语音消息] 格式: {format}"
            elif msg_type == 'video':
                media_id = root.find('MediaId').text
                thumb_media_id = root.find('ThumbMediaId').text
                content = f"[视频消息]"
            elif msg_type == 'location':
                location_x = root.find('Location_X').text
                location_y = root.find('Location_Y').text
                scale = root.find('Scale').text
                label = root.find('Label').text
                content = f"[位置消息] {label} ({location_x}, {location_y})"
            else:
                content = f"[{msg_type}类型消息]"
            
            return {
                'msg_type': msg_type,
                'from_user': from_user,
                'to_user': to_user,
                'create_time': create_time,
                'msg_id': msg_id,
                'content': content
            }
        except Exception as e:
            print(f"解析XML消息失败: {str(e)}")
            return None

# 示例用法
if __name__ == "__main__":
    # 这里填入你的Token、EncodingAESKey、CorpID
    token = "Sv3mVW1LOIU"
    encoding_aes_key = "NgHrjQdRw46FQIyTedXxJu6GrjqcoewHWlCQVLA76Ai"
    corp_id = "wxc4c1ad1d26dc0daf"

    wx_crypto = WXBizMsgCrypt(token, encoding_aes_key, corp_id)
    
    # 测试加解密
    encrypted = wx_crypto.encrypt("hello world")
    print("加密后：", encrypted)
    decrypted = wx_crypto.decrypt(encrypted)
    print("解密后：", decrypted)
    
    # 测试签名
    timestamp = str(int(time.time()))
    nonce = ''.join(random.sample(string.ascii_letters + string.digits, 16))
    signature = wx_crypto.get_signature(timestamp, nonce, encrypted)
    print("签名：", signature)
